import os
from ultralytics import YOLO
from matplotlib import rcParams
from matplotlib.font_manager import FontProperties

# Handle potential library conflicts
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# Set up Chinese font support for matplotlib
font_path = './simhei.ttf'
if not os.path.exists(font_path):
    print("Error: SimHei font file not found at the specified path.")
else:
    prop = FontProperties(fname=font_path)
    rcParams['font.family'] = prop.get_name()

def main():
    # Model path and dataset file
    model_path = './best.pt'
    data_path = './datasets/data.yaml'

    # Load YOLO model
    model = YOLO(model_path)

    # Train the model with enhanced preprocessing for lighting and partial target recognition
    results = model.train(
        data=data_path,
        epochs=3000,
        imgsz=640,
        batch=8,
        name='exp',
        workers=4,
        verbose=True,
        augment=True,          # Enable data augmentation
        mosaic=0.3,            # Reduced mosaic to focus on full target
        degrees=45.0,          # rotation range
        scale=1.0,             # Scale adjustment for diverse target sizes
        hsv_v=0.8,             # Increase brightness adjustment (HSV value channel)
        hsv_s=0.9,             # Increase saturation adjustment
        hsv_h=0.5,             # Moderate hue shift for color diversity
        fliplr=0.5,            # Horizontal flip to enhance diversity
        flipud=0.1,            # Minor vertical flip to simulate slight orientation variance
        patience=100,          # Early stopping patience to allow ample training
        lr0=0.01,              # Initial learning rate
        lrf=0.00001            # Final learning rate for gradual decay
    )

if __name__ == '__main__':
    main()
